import cv2
import numpy as np
from datetime import datetime
import math

# -------------------------- 1. 配置参数（需根据图像修改）--------------------------
IMAGE_PATH = "butterfly_with_ruler.jpeg"  # 带弧线的蝴蝶图像路径
BLUR_KERNEL = (3, 3)              # 高斯模糊核（去噪）
THRESHOLD_VALUE = 130             # 二值化阈值（调整至翅膀轮廓清晰）
EXPORT_TXT_PATH = "arc_length_result.txt"  # 结果导出路径

# 全局变量：存储鼠标选点、椭圆参数、弧长
arc_points = []
ellipse_params = None
arc_length = 0.0  # 弧长（像素单位）


# -------------------------- 2. 鼠标交互函数（选点+拟合椭圆）--------------------------
def mouse_callback(event, x, y, flags, param):
    global arc_points, ellipse_params, arc_length
    img_copy = param  # 传入图像副本，避免原图被修改

    # 左键点击：记录弧线上的点
    if event == cv2.EVENT_LBUTTONDOWN:
        arc_points.append((x, y))
        # 绘制选点（红色，带序号）
        cv2.circle(img_copy, (x, y), 5, (0, 0, 255), -1)
        cv2.putText(img_copy, str(len(arc_points)), (x + 10, y - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)
        print(f"已选点{len(arc_points)}: ({x}, {y})")

        # 选够5个点：拟合椭圆并计算弧长（点越多拟合越准）
        if len(arc_points) >= 5:
            # 转换为OpenCV要求的点格式（(n,1,2)的uint8数组）
            pts = np.array(arc_points, dtype=np.int32).reshape((-1, 1, 2))
            try:
                # 拟合椭圆（返回：(中心(x,y), (长轴,短轴), 旋转角)）
                ellipse_params = cv2.fitEllipse(pts)
                cv2.ellipse(img_copy, ellipse_params, (0, 255, 0), 2)  # 绘制拟合的椭圆

                # 计算弧长（取第一个点和最后一个点之间的弧）
                if len(arc_points) >= 2:
                    center = (ellipse_params[0][0], ellipse_params[0][1])
                    a = ellipse_params[1][0] / 2  # 长半轴
                    b = ellipse_params[1][1] / 2  # 短半轴
                    phi = np.radians(ellipse_params[2])  # 椭圆旋转角（转弧度）

                    # 求解起始点和结束点在椭圆上的参数角θ
                    theta_start = find_theta(center, arc_points[0], a, b, phi)
                    theta_end = find_theta(center, arc_points[-1], a, b, phi)

                    # 数值积分计算弧长
                    arc_length = ellipse_arc_length(a, b, theta_start, theta_end)
                    # 在图像上显示弧长
                    cv2.putText(img_copy, f"弧长: {arc_length:.1f} px",
                                (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 0, 255), 3)
                    print(f"计算完成！弧长: {arc_length:.1f} 像素")
            except:
                print("提示：选点不足或分布不合理，无法拟合椭圆！")


# -------------------------- 3. 辅助函数：求解椭圆参数角θ --------------------------
def find_theta(center, point, a, b, phi):
    """根据椭圆上的点，求解其对应的参数角θ（椭圆参数方程中的角度）"""
    dx = point[0] - center[0]
    dy = point[1] - center[1]
    # 将点旋转到椭圆的“原始坐标系”（抵消椭圆的旋转）
    x_rot = dx * math.cos(phi) + dy * math.sin(phi)
    y_rot = -dx * math.sin(phi) + dy * math.cos(phi)
    # 近似求解参数角θ（通过反正切，适合椭圆上的点）
    theta = math.atan2(y_rot * a, x_rot * b)
    return theta


# -------------------------- 4. 辅助函数：数值积分计算椭圆弧长 --------------------------
def ellipse_arc_length(a, b, theta_start, theta_end, num_segments=1000):
    """用辛普森法则数值积分，计算椭圆从θ_start到θ_end的弧长"""
    def integrand(theta):
        """椭圆参数方程的导数模长（弧长积分的被积函数）"""
        dx_dtheta = -a * math.sin(theta)
        dy_dtheta = b * math.cos(theta)
        return math.sqrt(dx_dtheta ** 2 + dy_dtheta ** 2)

    # 确保积分区间为正（若θ_end < θ_start，跨0则补2π）
    if theta_end < theta_start:
        theta_end += 2 * math.pi

    # 辛普森法则：将区间分成偶数段，用抛物线近似积分
    h = (theta_end - theta_start) / num_segments
    sum_odd, sum_even = 0.0, 0.0
    for i in range(1, num_segments):
        theta = theta_start + i * h
        if i % 2 == 1:
            sum_odd += integrand(theta)
        else:
            sum_even += integrand(theta)
    # 辛普森公式：(h/3) * [f(a) + f(b) + 4*奇数项和 + 2*偶数项和]
    integral = (h / 3) * (integrand(theta_start) + integrand(theta_end) + 4 * sum_odd + 2 * sum_even)
    return integral


# -------------------------- 5. 结果导出函数 --------------------------
def export_length_result(length, img_path, ellipse_info=None):
    """将弧线长度结果导出到TXT文件"""
    with open(EXPORT_TXT_PATH, "w", encoding="utf-8") as f:
        f.write(f"蝴蝶翅膀弧线长度测量报告\n")
        f.write(f"测量时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"图像路径：{img_path}\n")
        f.write(f"弧线长度（像素）：{length:.2f}\n")
        if ellipse_info:
            center, axes, angle = ellipse_info
            f.write(f"拟合椭圆中心：({center[0]:.1f}, {center[1]:.1f})\n")
            f.write(f"拟合椭圆长轴：{axes[0]:.1f}, 短轴：{axes[1]:.1f}\n")
            f.write(f"拟合椭圆旋转角：{angle:.1f} 度\n")
    print(f"\n结果已导出至：{EXPORT_TXT_PATH}")


# -------------------------- 6. 主流程执行 --------------------------
if __name__ == "__main__":
    # 步骤1：图像预处理（突出翅膀轮廓）
    img = cv2.imread(IMAGE_PATH)
    if img is None:
        print(f"错误：无法读取图像 {IMAGE_PATH}")
        exit(1)
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img_blur = cv2.GaussianBlur(img_gray, BLUR_KERNEL, 0)
    # 二值化（翅膀为白色，背景为黑色）
    _, img_bin = cv2.threshold(img_blur, THRESHOLD_VALUE, 255, cv2.THRESH_BINARY_INV)
    # 形态学膨胀（增强轮廓连续性）
    kernel = np.ones((2, 2), np.uint8)
    img_bin = cv2.dilate(img_bin, kernel, iterations=1)

    # 步骤2：创建图像副本用于绘制
    img_display = img.copy()
    # 显示操作提示
    cv2.putText(img_display, "操作提示：点击弧线上的点（至少5个）", (20, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)

    # 步骤3：绑定鼠标回调函数，开启交互选点
    cv2.namedWindow("Butterfly Wing Arc Length Measurement", cv2.WINDOW_NORMAL)
    cv2.setMouseCallback("Butterfly Wing Arc Length Measurement", mouse_callback, img_display)

    # 步骤4：等待用户操作（按ESC退出，按S保存结果）
    while True:
        cv2.imshow("Butterfly Wing Arc Length Measurement", img_display)
        key = cv2.waitKey(1) & 0xFF
        # 按ESC键：退出程序
        if key == 27:
            break
        # 按S键：导出结果（需先拟合椭圆）
        if key == ord('s'):
            if ellipse_params is not None:
                export_length_result(arc_length, IMAGE_PATH, ellipse_params)
            else:
                print("提示：请先选点拟合椭圆再导出结果！")

    # 释放资源
    cv2.destroyAllWindows()